package com.xuyanze.es.utils;

import com.alibaba.fastjson.JSON;
import com.xuyanze.es.EsLambdaQueryWrapper;
import com.xuyanze.es.enums.EsJoinType;
import com.xuyanze.es.enums.EsKeyWord;
import com.xuyanze.es.model.EsSegment;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.sort.SortBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;

import java.util.List;

/**
 * @author ：lzy
 * @date ：2023-02-10
 * @description：
 */
public class EsWrapperProcessor {

    private static final Logger LOGGER = LoggerFactory.getLogger(EsWrapperProcessor.class);

    private static boolean enableDsl = false;

    public static void setEnableDsl(boolean enableDsl) {
        EsWrapperProcessor.enableDsl = enableDsl;
    }

    public static NativeSearchQuery initNativeSearchQuery(EsLambdaQueryWrapper wrapper){
        BoolQueryBuilder boolQueryBuilder = initQueryBuilder(wrapper);
        NativeSearchQueryBuilder nativeSearchQueryBuilder = new NativeSearchQueryBuilder();
        QueryBuilder queryBuilder = null;

        // constScore
        if(wrapper.isConstScore()){
            queryBuilder = QueryBuilders.constantScoreQuery(boolQueryBuilder);
        }else{
            queryBuilder = boolQueryBuilder;
        }

        // 查询指定字段
        if(wrapper.getWithFields()!=null){
            nativeSearchQueryBuilder.withFields(wrapper.getWithFields());
        }

        // 排序
        if(wrapper.getPageable()!=null){
            nativeSearchQueryBuilder.withPageable(wrapper.getPageable());
        }

        // 排序字段
        List<SortBuilder> sortBuilderList = wrapper.getSortBuilderList();
        for (SortBuilder sortBuilder : sortBuilderList) {
            nativeSearchQueryBuilder.withSort(sortBuilder);
        }

        // 聚会函数
        if(wrapper.getAggregationBuilder()!=null){
            nativeSearchQueryBuilder.addAggregation(wrapper.getAggregationBuilder());
        }

        nativeSearchQueryBuilder.withQuery(queryBuilder);
        if(enableDsl) System.out.println(queryBuilder);;
        return nativeSearchQueryBuilder.build();
    }

    public static BoolQueryBuilder initQueryBuilder(EsLambdaQueryWrapper wrapper){
        List<EsSegment> segments = wrapper.getSegments();
        if(segments == null || segments.isEmpty()){
            return null;
        }
        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
        for (int i = 0; i < segments.size(); i++) {
            EsSegment segment = segments.get(i);
            EsKeyWord keyWord = segment.getKeyWord();
            if(keyWord == EsKeyWord.OR_JOIN){
                segments.get(i - 1).setJoinType(EsJoinType.SHOULD);
                segments.get(i + 1).setJoinType(EsJoinType.SHOULD);
                i++; continue;
            }
        }
        // 拼接query
        buildQueryBuilder(boolQueryBuilder,segments);
        // minimumShouldMatch
        if(wrapper.getMinimumShouldMatch()!=null) boolQueryBuilder.minimumShouldMatch(wrapper.getMinimumShouldMatch());
        return boolQueryBuilder;
    }

    private static void buildQueryBuilder(BoolQueryBuilder boolQueryBuilder, List<EsSegment> segments) {
        for (int i = 0; i < segments.size(); i++) {
            EsSegment segment = segments.get(i);
            addQuery(boolQueryBuilder,segment);
        }
    }

    private static void addQuery(BoolQueryBuilder boolQueryBuilder, EsSegment segment) {
        if(segment.bindQuery() == null) return;
        switch (segment.getJoinType()){
           case MUST:
               boolQueryBuilder.must(segment.bindQuery());
               break;
           case MUST_NOT:
               boolQueryBuilder.mustNot(segment.bindQuery());
               break;
           case SHOULD:
               boolQueryBuilder.should(segment.bindQuery());
               break;
           case FILTER:
               boolQueryBuilder.filter(segment.bindQuery());
               break;
           default:
               throw new RuntimeException("不支持的JoinType"+segment.getJoinType().name());
       }
    }
}
